# from langchain.memory import ConversationSummaryMemory
# from qw_model import qw_model
#
# model = qw_model()
# memory = ConversationSummaryMemory(llm=model)
# memory.save_context({"input": "你好"}, {"output": "有什么事"})
# memory.save_context({"input": "我正在为聊天机器人撰写更好的文档"}, {"output": "噢，听起来很辛苦"})
# memory.save_context({"input": "是的，但是这很值得努力"}, {"output": "同意，好的文档很重要！"})
# print(memory.load_memory_variables({}))
#---------------------------------------------------------------------
# from typing import Literal
# from langchain_core.messages import SystemMessage, RemoveMessage, HumanMessage
# from langgraph.checkpoint.memory import MemorySaver
# from langgraph.graph import MessagesState, StateGraph, START, END
#
# memory = MemorySaver()
#
#
# # We will add a `summary` attribute (in addition to `messages` key,
# # which MessagesState already has)
# class State(MessagesState):
#     summary: str
#
#
# # We will use this model for both the conversation and the summarization
# model = qw_model()
#
#
# # Define the logic to call the model
# def call_model(state: State):
#     # If a summary exists, we add this in as a system message
#     summary = state.get("summary", "")
#     if summary:
#         system_message = f"Summary of conversation earlier: {summary}"
#         messages = [SystemMessage(content=system_message)] + state["messages"]
#     else:
#         messages = state["messages"]
#     response = model.invoke(messages)
#     # We return a list, because this will get added to the existing list
#     return {"messages": [response]}
#
#
# # We now define the logic for determining whether to end or summarize the conversation
# def should_continue(state: State) -> Literal["summarize_conversation", END]:
#     """Return the next node to execute."""
#     messages = state["messages"]
#     # If there are more than six messages, then we summarize the conversation
#     if len(messages) > 6:
#         return "summarize_conversation"
#     # Otherwise we can just end
#     return END
#
#
# def summarize_conversation(state: State):
#     # First, we summarize the conversation
#     summary = state.get("summary", "")
#     if summary:
#         # If a summary already exists, we use a different system prompt
#         # to summarize it than if one didn't
#         summary_message = (
#             f"This is summary of the conversation to date: {summary}\n\n"
#             "Extend the summary by taking into account the new messages above:"
#         )
#     else:
#         summary_message = "Create a summary of the conversation above:"
#
#     messages = state["messages"] + [HumanMessage(content=summary_message)]
#     response = model.invoke(messages)
#     # We now need to delete messages that we no longer want to show up
#     # I will delete all but the last two messages, but you can change this
#     delete_messages = [RemoveMessage(id=m.id) for m in state["messages"][:-2]]
#     return {"summary": response.content, "messages": delete_messages}
#
#
# # Define a new graph
# workflow = StateGraph(State)
#
# # Define the conversation node and the summarize node
# workflow.add_node("conversation", call_model)
# workflow.add_node(summarize_conversation)
#
# # Set the entrypoint as conversation
# workflow.add_edge(START, "conversation")
#
# # We now add a conditional edge
# workflow.add_conditional_edges(
#     # First, we define the start node. We use `conversation`.
#     # This means these are the edges taken after the `conversation` node is called.
#     "conversation",
#     # Next, we pass in the function that will determine which node is called next.
#     should_continue,
# )
#
# # We now add a normal edge from `summarize_conversation` to END.
# # This means that after `summarize_conversation` is called, we end.
# workflow.add_edge("summarize_conversation", END)
#
# # Finally, we compile it!
# app = workflow.compile(checkpointer=memory)

#-------------------------------------------------------------
import bs4
from langchain.chains import create_history_aware_retriever, create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_chroma import Chroma
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_community.document_loaders import WebBaseLoader
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_community.llms import Tongyi
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

llm = Tongyi()

### 构造检索器 ###
bs4_strainer = bs4.SoupStrainer(class_=("title-article", "baidu_pl"))
loader = WebBaseLoader(
    web_paths=("https://blog.csdn.net/fengshi_fengshi/article/details/144359686",),
    bs_kwargs={"parse_only": bs4_strainer},
)
docs = loader.load()

text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
splits = text_splitter.split_documents(docs)
vectorstore = Chroma.from_documents(documents=splits, embedding=DashScopeEmbeddings())
retriever = vectorstore.as_retriever()

### 情境化问题 ###
contextualize_q_system_prompt = (
    "给定聊天记录和最新用户问题 "
    "这可能会引用聊天历史中的上下文， "
    "制定一个可以理解的独立问题 "
    "没有聊天记录。不要回答问题， "
    "如果需要，只需重新规划，否则按原样退回。"
)
contextualize_q_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", contextualize_q_system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ]
)
history_aware_retriever = create_history_aware_retriever(
    llm, retriever, contextualize_q_prompt
)

### 回答问题 ###
system_prompt = (
    "你是问答任务的助理。 "
    "使用以下检索到的上下文来回答问题。 "
    "如果你不知道答案，就说你不知道。 "
    "最多使用五句话，并保持回答简明扼要。 "
    "\n\n"
    "{context}"
)
qa_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ]
)
question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)

rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)

### 有状态地管理聊天记录 ###
store = {}

def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = ChatMessageHistory()
    return store[session_id]


conversational_rag_chain = RunnableWithMessageHistory(
    rag_chain,
    get_session_history,
    input_messages_key="input",
    history_messages_key="chat_history",
    output_messages_key="answer",
)




